from geometry_msgs.msg import TransformStamped
from sensor_msgs.msg import Range
import rclpy
from rclpy.node import Node
import math
import time
import message_filters
from tf2_ros import TransformBroadcaster  
import tf_transformations
from turtlesim.msg import Pose


class FramePublisher(Node):

    def __init__(self):
        super().__init__('cranebot_frame_publisher')
        self.imu_sub = message_filters.Subscriber(self, Imu, '/cranebot/imu')
        self.tof1_sub = message_filters.Subscriber(self, Range, '/cranebot/tof1')
        self.tof2_sub = message_filters.Subscriber(self, Range, '/cranebot/tof2')
        ts = message_filters.TimeSynchronizer([self.imu_sub, self.tof1_sub, self.tof2_sub], 10)
        ts.registerCallback(self.handle_cranebot_pose)
        self.br = TransformBroadcaster(self)
        

        self.hook_distance_init = 20/100
        self.arm_distance_init = 7/100

    def handle_cranebot_pose(self, imu, tof1, tof2):
        self.get_logger().info('tof1'+str(tof1.range)+'tof2'+str(tof2.range))
        t0 = TransformStamped()

        ## Tof is used to measure the movement of the cable wire and the arm
        self.arm_move = float(tof1.range)
        self.wire_move = float(tof2.range)
        self.hook_distance = self.hook_distance_init + self.wire_move/4 - self.arm_move/4
        self.arm_distance = self.arm_distance_init + self.arm_move

        #self.get_logger().info(str([hook_distance, arm_distance]))
        # Read message content and assign it to
        # corresponding tf variables

        t0.header.stamp = self.get_clock().now().to_msg()
        t0.header.frame_id = 'world'
        t0.child_frame_id = 'car'

        # Turtle only exists in 2D, thus we get x and y translation
        # coordinates from the message and set the z coordinate to 0
        t0.transform.translation.x = 0.0
        t0.transform.translation.y = 0.0
        t0.transform.translation.z = 0.0

        # For the same reason, turtle can only rotate around one axis
        # and this why we set rotation in x and y to 0 and obtain
        # rotation in z axis from the message
        q = tf_transformations.quaternion_from_euler(0, 0, 0)
        t0.transform.rotation.x = q[0]
        t0.transform.rotation.y = q[1]
        t0.transform.rotation.z = q[2]
        t0.transform.rotation.w = q[3]

        t1 = TransformStamped()
        t1.header.stamp = self.get_clock().now().to_msg()
        t1.header.frame_id = 'car'
        t1.child_frame_id = 'cab'

        # Turtle only exists in 2D, thus we get x and y translation
        # coordinates from the message and set the z coordinate to 0
        t1.transform.translation.x = 0.0
        t1.transform.translation.y = 0.0
        t1.transform.translation.z = 0.0

        t2 = TransformStamped()
        t2.header.stamp = self.get_clock().now().to_msg()
        t2.header.frame_id = 'cab'
        t2.child_frame_id = 'arm_throttle'
        t2.transform.translation.x = 0.0
        t2.transform.translation.y = 0.0
        t2.transform.translation.z = 0.22    
        q2 = tf_transformations.quaternion_from_euler(0, -30/180*math.pi, 0)
        t2.transform.rotation.x = q2[0]
        t2.transform.rotation.y = q2[1]
        t2.transform.rotation.z = q2[2]
        t2.transform.rotation.w = q2[3]      

        t3 = TransformStamped()
        t3.header.stamp = self.get_clock().now().to_msg()
        t3.header.frame_id = 'arm_throttle'
        t3.child_frame_id = 'tof'
        t3.transform.translation.x = 1.5
        t3.transform.translation.y = 0.0
        t3.transform.translation.z = 0.0    
        q3 = tf_transformations.quaternion_from_euler(0, 0, 0)
        t3.transform.rotation.x = q3[0]
        t3.transform.rotation.y = q3[1]
        t3.transform.rotation.z = q3[2]
        t3.transform.rotation.w = q3[3]

        t4 = TransformStamped()
        t4.header.stamp = self.get_clock().now().to_msg()
        t4.header.frame_id = 'tof'
        t4.child_frame_id = 'arm_end'
        t4.transform.translation.x = self.arm_distance
        t4.transform.translation.y = 0.0
        t4.transform.translation.z = 0.0    
        q4 = tf_transformations.quaternion_from_euler(0, 0, 0)
        t4.transform.rotation.x = q4[0]
        t4.transform.rotation.y = q4[1]
        t4.transform.rotation.z = q4[2]
        t4.transform.rotation.w = q4[3]    

        t5 = TransformStamped()
        t5.header.stamp = self.get_clock().now().to_msg()
        t5.header.frame_id = 'arm_end'
        t5.child_frame_id = 'hook'
        l = self.hook_distance
        t5.transform.translation.x = l * math.sin(-30/180*math.pi)
        t5.transform.translation.y = 0.0 
        t5.transform.translation.z = -l * math.cos(-30/180*math.pi)  
        q5 = tf_transformations.quaternion_from_euler(0, 30/180*math.pi, 0)
        t5.transform.rotation.x = q5[0]
        t5.transform.rotation.y = q5[1]
        t5.transform.rotation.z = q5[2]
        t5.transform.rotation.w = q5[3]            


        # Send the transformation
        self.br.sendTransform(t0)
        self.br.sendTransform(t1)
        self.br.sendTransform(t2)
        self.br.sendTransform(t3)
        self.br.sendTransform(t4)
        self.br.sendTransform(t5)

def main():
    rclpy.init()
    node = FramePublisher()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass

    rclpy.shutdown()

if __name__ == '__main__':
    main()